# 功能：特征构造 + 特征变换 + 特征选择 + 数据质量报告 + 特征重要性可视化

import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
import os
import matplotlib.pyplot as plt
import seaborn as sns


def create_new_features(df):
    """
    构造新特征
    :param df: 预处理后的DataFrame
    :return: 添加新特征后的DataFrame
    """
    print("【开始】特征构造...")
    df['fare_density'] = df['Average_Fare'] / df['NonStopMiles']
    print(" 新增特征 fare_density")

    #  新增特征：票价密度分段（特征变换 - 离散化）
    bins = [0, 0.5, 1.0, 2.0, float('inf')]
    labels = ['Very Low', 'Low', 'Medium', 'High']
    df['fare_density_category'] = pd.cut(df['fare_density'], bins=bins, labels=labels)

    print(" 新增离散化特征 fare_density_category")
    return df


def plot_feature_importance(importance, output_dir):
    """绘制特征重要性条形图"""
    plt.figure(figsize=(12, 8))

    # 筛选重要性大于阈值的特征
    importance_filtered = importance[importance > 0.01]

    # 如果没有特征满足条件，使用前10个特征
    if importance_filtered.empty:
        importance_filtered = importance.head(10)

    importance_filtered.sort_values().plot(kind='barh')
    plt.title("Top Feature Importances")
    plt.xlabel("Importance")
    plt.ylabel("Features")
    plt.yticks(rotation=0)  # 旋转Y轴标签

    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "feature_importance.png"))
    plt.close()
def select_important_features(X, y):
    print("【开始】特征选择...")
    model = RandomForestRegressor(n_estimators=100, random_state=42)
    model.fit(X, y)

    importance = pd.Series(model.feature_importances_, index=X.columns).sort_values(ascending=False)
    top_features = importance.head(50).index.tolist()
    print(" 筛选出前50个重要特征")

    # 可视化特征重要性
    os.makedirs("../data/output", exist_ok=True)
    plot_feature_importance(importance, "../data/output")

    return top_features


def generate_quality_report(df):
    """
    输出数据质量报告
    :param df: DataFrame
    """
    print("\n 特征工程后数据质量报告：")
    missing_summary = df.isnull().sum()
    missing_rate = (missing_summary / df.shape[0]) * 100
    quality_report = pd.DataFrame({
        'Missing Count': missing_summary,
        'Missing Rate (%)': missing_rate,
        'Unique Count': df.nunique(),
        'Data Type': df.dtypes
    })
    print(quality_report)


if __name__ == "__main__":
    df = pd.read_parquet("../data/preprocessed/cleaned_data.parquet")
    df_with_new_features = create_new_features(df.copy())

    # 新增 One-Hot 编码
    print(" 正在进行 One-Hot 编码...")
    df_encoded = pd.get_dummies(df_with_new_features, columns=['fare_density_category'], drop_first=True)

    # 假设目标变量是 'Average_Fare'
    X = df_encoded.drop(columns=['Average_Fare'])
    y = df_encoded['Average_Fare']

    selected_features = select_important_features(X, y)
    final_df = df_encoded[selected_features + ['Average_Fare']]

    #  输出数据质量报告
    generate_quality_report(final_df)

    #  输出路径
    output_path = "../data/features/feature_engineered_data.parquet"

    #  自动创建父目录（如果不存在）
    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    #  保存 Parquet 文件
    final_df.to_parquet(output_path)
    print(f" 已保存特征工程结果到 {output_path}")
